/*
 * fdm2dutils.h
 *
 *  Created on: 12-Apr-2010
 *      Author: nish
 */

#ifndef FDM2DUTILS_H_
#define FDM2DUTILS_H_

#include "fdmconstants.h"

enum POT_TYPE{
	STEP, SQUARE_PULSE, GRADIENT,GRADIENT_STEP, COLOUMBIC
};

double BLC = 0;
double BRC = 1;
double TLC = 2;
double TRC = 3;
double BEN = 4;
double LEN = 5;
double REN = 6;
double TEN = 7;
double MID = 8;

int m;
int n;
double hmat_size;
double xlength=3e-9;
double ylength=3e-9;
double ax=1e-10;
double ay=1e-10;
double t0x = 0;
double t0y = 0;

void init2DConstants();
gsl_matrix* form2DHMatrix();
gsl_vector* markNodes();
void save2DHMatrix(gsl_matrix* H, char* file);
void solve2DEig(gsl_matrix* H, gsl_vector* eig_val,gsl_matrix* eig_vect);
gsl_matrix* form2DPotMatrix();

void init2DConstants(){
	m = xlength/ax + 1;
	n = ylength/ay + 1;
	hmat_size = m*n;
	t0x = pow(hbar,2)/(2*e_mass*pow(ax,2)*e_charge);
	t0y = pow(hbar,2)/(2*e_mass*pow(ay,2)*e_charge);

	printf("t0x = %f && t0y = %f\n",t0x,t0y);
}

gsl_vector* markNodes(){
	gsl_vector* node_marker = gsl_vector_alloc(hmat_size);
	int i = 0;

	for(i=0;i<hmat_size;i++){
		if(i==0){ // Bottom left corner node
			gsl_vector_set(node_marker,i,BLC);
		}else if(i/n == 0){
			if((i+1)%n == 0){ //bottom right corner node
				gsl_vector_set(node_marker,i,BRC);
			}else{//add && i%n != 0 //bottom edge node
				gsl_vector_set(node_marker,i,BEN);
			}
		}else if(i%n ==0){
			//left edge, including the top left node
			if(i/n == m-1){
				//top left corner node
				gsl_vector_set(node_marker,i,TLC);
			}else{
				//left edge node
				gsl_vector_set(node_marker,i,LEN);
			}
		}else if((i+1)%n ==0){
			//right edge including the top right corner node
			if((i+1)/n == m){
				//top right corner node
				gsl_vector_set(node_marker,i,TRC);
			}else{
				//right edge node
				gsl_vector_set(node_marker,i,REN);
			}
		}else if (i/n == m-1){ //add && i %n != 0
			// top edge node
			gsl_vector_set(node_marker,i,TEN);
		}else{
			//middle node
			gsl_vector_set(node_marker,i,MID);
		}
	}


	return node_marker;
}

gsl_matrix* form2DHMatrix(){
	gsl_matrix* H = gsl_matrix_alloc(hmat_size,hmat_size);
	gsl_vector* node_marker = gsl_vector_alloc(hmat_size);
	int i=0;
	int j=0;
	double node = 0;
	printf("marking the nodes in Hmatrix\n");
	node_marker = markNodes();

//	for(i=0;i<hmat_size;i++){
//		printf("The node %d is marked as %f\n",i,gsl_vector_get(node_marker,i));
//	}

	printf("building HMatrix\n");
	// iterate through all elements of H -> hmat_size * hmat_size - set the type of nodes (
	for(i=0;i<hmat_size;i++){

		node = gsl_vector_get(node_marker,i);
		if( node == BLC){
			gsl_matrix_set(H,i,i,t0x+t0y);
			gsl_matrix_set(H,i,i+n,-t0y);
			gsl_matrix_set(H,i,i+1,-t0x);
		}else if(node == BEN){
			// periodic in y diretion
			gsl_matrix_set(H,i,i,2*t0x+2*t0y);
			gsl_matrix_set(H,i,i+n,-t0y);
			gsl_matrix_set(H,i,i+1,-t0x);
			gsl_matrix_set(H,i,i-1,-t0x);

			gsl_matrix_set(H,i,n-1,-t0y);
		}else if(node == BRC){
			gsl_matrix_set(H,i,i,t0x+t0y);
			gsl_matrix_set(H,i,i+n,-t0y);
			gsl_matrix_set(H,i,i-1,-t0x);
		}else if(node == REN){
			// periodic in x direction
			gsl_matrix_set(H,i,i,2*t0x+2*t0y);
			gsl_matrix_set(H,i,i+n,-t0y);
			gsl_matrix_set(H,i,i-n,-t0y);
			gsl_matrix_set(H,i,i-1,-t0x);

			gsl_matrix_set(H,i,0,-t0x);
		}else if(node == LEN){
			// periodic in x direction
			gsl_matrix_set(H,i,i,2*t0x+2*t0y);
			gsl_matrix_set(H,i,i+n,-t0y);
			gsl_matrix_set(H,i,i-n,-t0y);
			gsl_matrix_set(H,i,i+1,-t0x);

			gsl_matrix_set(H,i,n-1,-t0x);
		}else if (node == TEN){
			//periodic in y direction
			gsl_matrix_set(H,i,i,2*t0x+2*t0y);
			gsl_matrix_set(H,i,i+1,-t0x);
			gsl_matrix_set(H,i,i-n,-t0y);
			gsl_matrix_set(H,i,i-1,-t0x);

			gsl_matrix_set(H,i,0,-t0y);
		}else if(node == TLC){
			gsl_matrix_set(H,i,i,t0x+t0y);
			gsl_matrix_set(H,i,i-n,-t0y);
			gsl_matrix_set(H,i,i+1,-t0x);
		}else if(node == TRC){
			gsl_matrix_set(H,i,i,t0x+t0y);
			gsl_matrix_set(H,i,i-n,-t0y);
			gsl_matrix_set(H,i,i-1,-t0x);
		}else if(node == MID){
			gsl_matrix_set(H,i,i,2*t0x+2*t0y);
			gsl_matrix_set(H,i,i+1,-t0x);
			gsl_matrix_set(H,i,i-n,-t0y);
			gsl_matrix_set(H,i,i-1,-t0x);
			gsl_matrix_set(H,i,i+n,-t0y);
		}
	}
	return H;
}


void save2DHMatrix(gsl_matrix* H, char* file){
	FILE* f = fopen(file,"w");
	gsl_matrix_fprintf(f,H,"%.5g");
}

void solve2DEig(gsl_matrix* H, gsl_vector* eig_val,gsl_matrix* eig_vect){

	//initialize the eigen vect workspace and get solve for eigen values and vectors:
	gsl_eigen_symmv_workspace* workspace = gsl_eigen_symmv_alloc(hmat_size);
	gsl_eigen_symmv(H,eig_val,eig_vect,workspace);
	gsl_eigen_symmv_free(workspace);
	gsl_eigen_symmv_sort(eig_val,eig_vect,GSL_EIGEN_SORT_ABS_ASC);
	FILE* f = fopen("2DeigVal.dat","w");
	gsl_vector_fprintf(f,eig_val,"%.5g");
	fclose(f);
	f = fopen("2DeigVet.dat","w");
	gsl_matrix_fprintf(f,eig_vect,"%.5g");
	fclose(f);

	return;
}


void saveEigMatrices(gsl_matrix* mat, int eigNum){
	FILE* f = fopen("2DEigTest.dat","w");
// is it really useful to convert the eigen vect to a matrix, don't we just need to see it visually now.
//	int imat,jmat;
	int x,y;
	int i=0;
	int j=0;

	x=-1;y=0;

	for(i=0;i<hmat_size;i++){
		if((i) % (n) == 0){ // if reaching the end of the column, then go to the next row.
			x++;
			y=0;
			fprintf(f,"\n");
		}
		fprintf(f,"%d\t%d\t%f\n",x,y++,pow(gsl_matrix_get(mat,eigNum,i),2)); //printing in the gnu plot required matrix format for 3d splot
	}
	fclose(f);

}


gsl_matrix* form2DPotMatrix(enum POT_TYPE potential_type){
	gsl_matrix* pot = gsl_matrix_alloc(m,n);
//	gsl_matrix_set_zero(pot);

	int i,j;
	switch(potential_type){
		case STEP:
			for(i=0;i<m;i++){
				for(j=0;j<n;j++){
					if(j>=n/2){
						gsl_matrix_set(pot,i,j,10);
					}else{
						gsl_matrix_set(pot,i,j,0);
					}
				}
				printf("\n");
			}
			break;

		default:
			break;
	}

	return pot;
}

void savePotMatrix(gsl_matrix* U){
	int i=0;
	int j=0;
	FILE* f = fopen("potMat.dat","w");
	double d;
	for(i=0;i<m;i++){
		for(j=0;j<n;j++){
			fprintf(f,"%d\t%d\t%.5g\n",i,j,gsl_matrix_get(U,i,j));
		}
		fprintf(f,"\n");
	}
	fclose(f);
}

/*
 * 2D matrix - random generated data for gnuplot testing.
 * The vector contains z data for given x and y. Have to split the
 */
void form2DMatrix(){
	int m = 3;
	int n = 5;
	int nTotal = m*n;
	int i=0;
	int j=0;
	srand(time(NULL));

	int x,y;

	gsl_vector* vect = gsl_vector_alloc(nTotal);
	for(i=0;i<nTotal;i++){
		gsl_vector_set(vect,i,rand()%10);
	}

	for(i=0;i<nTotal;i++){
		printf("vect %d -> %f\n",i,gsl_vector_get(vect, i));
	}

	FILE* f = fopen("2DTest.dat","w");
//	gsl_matrix mat = gsl_matrix_alloc(m,n);
// is it really useful to convert the eigen vect to a matrix, don't we just need to see it visually now.
//	int imat,jmat;
	x=0;y=0;

	for(i=0;i<nTotal;i++){
		printf("\n>>%d",(i)%(n));
		if((i) % (n) == 0){ // if reaching the end of the column, then go to the next row.
			x++;
			y=0;
			fprintf(f,"\n");
		}
		fprintf(f,"%d\t%d\t%f\n",x,y++,gsl_vector_get(vect,i)); //printing in the gnu plot required matrix format for 3d splot
	}
	fclose(f);
}
#endif /* FDM2DUTILS_H_ */
